from androguard.misc import AnalyzeAPK
import numpy as np
import pickle
import random
random.seed(0)


def get_intent(a):

    xml = a.get_android_manifest_axml().get_xml()
    str_xml = str(xml)
    xmls = str_xml.split(">")
    # print(xml)
    flag = False
    intent = []
    for i in range(0, len(xmls)):
        if "<intent-filter/" in xmls[i]:
            flag = False
        elif "</intent-filter" in xmls[i]:
            flag = False
        elif "<intent-filter" in xmls[i]:
            flag = True
        try:
            if flag and "</intent-filter" not in xmls[i+1]:
                intent.append(xmls[i+1])
        except:
            print('异常')
    intents = []
    for item in intent:
        names = item.split('"')
        if len(names) > 1:
            intents.append(names[1])

    return intents


def get_apis(dx):

    cs = []
    for i in dx.get_classes():
        classname = str(i.get_class())
        if classname.startswith('L'):
            internalclass = classname.split(
                "->L")[1].replace('/', '.').replace(";", "")
            cs.append(internalclass)

    apis = []
    for m in dx.get_methods():
        methodname = str(m.get_method())
        if methodname[1:].split(";->")[0].replace('/', '.') not in cs:
            api = methodname.replace("/", ".").replace(";->", ".")
            api = api[1:].split("(")[0].replace("<", "").replace(">", "")
            apis.append(api)
    return list(set(apis))


def get_permission(a):
    return a.get_permissions()


def get_select_features(file_path):
    api = []
    intent = []
    permission = []
    with open(file_path, 'r') as f_in:
        line = f_in.readline().strip()
        while line.count('\n') != len(line):
            intent.append(line)
            line = f_in.readline().strip()
        line = f_in.readline().strip()
        while line.count('\n') != len(line):
            permission.append(line)
            line = f_in.readline().strip()
        line = f_in.readline().strip()
        while line:
            api.append(line)
            line = f_in.readline().strip()
    return intent, permission, api


def analyse(file, raw):
    a, d, dx = AnalyzeAPK(file, raw=raw)
    intents = get_intent(a)
    apis = get_apis(dx)
    permission = get_permission(a)
    return intents, permission, apis


def extract_features(input_features, select_features):
    select_INTENTS, select_PERMISSION, select_API = select_features
    intents, permission, apis = input_features
    len_INTENTS = len(select_INTENTS)
    len_PERMISSION = len(select_PERMISSION)
    len_API = len(select_API)
    features = [0] * (len_INTENTS + len_PERMISSION + len_API)
    for item in list(set(intents)):
        try:
            index = select_INTENTS.index(item)
            features[index] = 1
        except:
            continue
    for item in list(set(permission)):
        try:
            index = select_PERMISSION.index(item)
            features[len_INTENTS + index] = 1
        except:
            continue
    for item in list(set(apis)):
        try:
            index = select_API.index(item)
            features[len_INTENTS + len_PERMISSION + index] = 1
        except:
            continue
    return features


def process(file, select_features):
    features_string = analyse(file, raw=True)
    features_vec = extract_features(features_string, select_features)
    features_vec_np = np.array([features_vec])
    pred = model.predict(features_vec_np)
    return pred[0], features_string, features_vec


def process_string(features_string, select_features):
    features_vec = extract_features(features_string, select_features)
    res = model_predict(features_vec)
    return res, features_string, features_vec


def model_predict(features_vec):
    features_vec_np = np.array([features_vec])
    pred = model.predict(features_vec_np)
    return pred[0]


model = pickle.load(open('./static/model/android/android_model.pkl', 'rb'))
select_features_path = "./static/model/android/android_features.txt"
select_features = get_select_features(select_features_path)

if __name__ == '__main__':
    file_path = "../TestAPK/apk_恶意/com.feamber.clandefenders.apk"
    with open(file_path, "rb") as f_in:
        file = f_in.read()

    res = process(file, select_features)
    print(res)
